import os, sys
import re, time
import subprocess
import multiprocessing
from datetime import datetime


class Info():

    def __init__(self, current_line, next_line):
        current_elements = [element.strip() for element in re.split(r'\||\s+|/', current_line) if element.strip()]
        next_elements = [element.strip() for element in re.split(r'\||\s+|/', next_line) if element.strip()]
        if len(current_elements) == 7 and re.match(r"\d", current_elements[0]):
            self.flag = True
            self.npu = current_elements[0]
            self.name = current_elements[1]
            self.health = current_elements[2]
            self.power = current_elements[3]
            self.temp = current_elements[4]
            self.hugepages = current_elements[5]
            self.hugepages_total = current_elements[6]
            self.chip = next_elements[0]
            self.phyid = next_elements[1]
            self.busid = next_elements[2]
            self.aicore = next_elements[3]
            self.memuse = next_elements[4]
            self.memuse_total = next_elements[5]
            self.hbm = next_elements[6]
            self.hbm_total = next_elements[7]
        else:
            self.flag = False


    def __str__(self):
        # ID, Power, Temp, AICore, HBM
        return f"{self.phyid},{self.power},{self.temp},{self.aicore},{self.hbm}"


def monitor(t0):
    index = 0
    infos = []
    now = datetime.now()
    lines = subprocess.run("npu-smi info", shell=True, text=True, capture_output=True).stdout.split('\n')
    while index < len(lines) - 1:
        info = Info(lines[index], lines[index + 1])
        index += 2 if info.flag else 1
        if info.flag:
            infos.append(str(info))
            if info.phyid == "0":
                print(f"{now} >> AICore: {info.aicore}")
    dt = (now - t0).total_seconds()
    return f"{now},{dt}," + ",".join(infos)

def run(cmd):
    time.sleep(2)
    subprocess.run(cmd, shell=True, text=True, capture_output=True)
    time.sleep(3)

# 处理输入的监控命令
t0 = datetime.now()
name = sys.argv[1] if len(sys.argv) > 1 else "hbm"
cmd = sys.argv[2] if len(sys.argv) > 2 else "./bin/hbm"

# 启动被监控的命令
print(f"name: {name}, cmd: {cmd}")
thread = multiprocessing.Process(target=run, args=(cmd,))
thread.start()

# 保存结果
if not os.path.exists("log"):
    os.mkdir("log")
with open(f"log/monitor-{name}.csv", 'w+') as file:
    line = "time,dt_s"
    for _ in range(8):
        line += ",phy-id,power_W,temp_C,aicore_%,hbm_MB"
    file.writelines(f"{line}\n")
    while thread.is_alive():
        file.writelines(f"{monitor(t0)}\n")
        file.flush()
        if (datetime.now() - t0).total_seconds() > 300:
            thread.terminate()
